import json, jieba
import numpy as np
from sentence_transformers import SentenceTransformer
from sklearn.cluster import KMeans
from sklearn.feature_extraction.text import TfidfVectorizer

questions = json.load(open("questions.json"))
print(questions[0])

model = SentenceTransformer('../hugging-face-model/BAAI/bge-small-zh-v1.5/')
question_sentences = [x['question'] for x in questions]
question_embeddings = model.encode(question_sentences, normalize_embeddings=True)

km = KMeans(n_clusters=20)
km.fit(question_embeddings)

print("bge编码--->")
for label in np.unique(km.labels_):
    cluster_idxs = np.where(km.labels_ == label)[0][:2]
    cluster_text = [questions[x]["question"] for x in cluster_idxs]

    print("聚类类别：", label)
    print(cluster_text)
    print("")

# 对文本进行分词
question_words = [' '.join(jieba.lcut(x['question'])) for x in questions]

# 提取TFIDF
tfidf = TfidfVectorizer()
tfidf.fit(question_words)
question_feat = tfidf.transform(question_words)

km = KMeans(n_clusters=20)
km.fit(question_feat)

print("TFIDF--->")
for label in np.unique(km.labels_):
    cluster_idxs = np.where(km.labels_ == label)[0][:2]
    cluster_text = [questions[x]["question"] for x in cluster_idxs]

    print("聚类类别：", label)
    print(cluster_text)
    print("")